#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
RAG系统调试脚本
诊断和修复chat方法的返回值问题
"""

import traceback
from tools.RAGChatBot import RAGChatBot


def test_rag_initialization():
    """测试RAG系统初始化"""
    print("🔧 测试RAG系统初始化...")
    
    try:
        bot = RAGChatBot(
            ollama_url="http://172.26.32.1:11434",
            milvus_collection="workspace_files",
            default_model="deepseek-r1:14b"
        )
        
        print(f"✅ RAG机器人创建成功")
        print(f"   Ollama URL: {bot.ollama_client.base_url}")
        print(f"   当前模型: {bot.current_model}")
        print(f"   Milvus集合: {bot.search_tool.collection_name}")
        
        # 测试初始化
        print("\n🚀 开始初始化...")
        success = bot.initialize()
        
        if success:
            print("✅ 初始化成功")
            return bot
        else:
            print("❌ 初始化失败")
            return None
            
    except Exception as e:
        print(f"❌ 创建RAG机器人失败: {e}")
        print("详细错误信息:")
        traceback.print_exc()
        return None


def test_ollama_connection(bot):
    """测试Ollama连接"""
    if not bot:
        return False
        
    print("\n🔌 测试Ollama连接...")
    
    try:
        # 测试连接
        if bot.ollama_client.test_connection():
            print("✅ Ollama连接正常")
            
            # 获取模型列表
            models = bot.ollama_client.list_models()
            print(f"📋 找到 {len(models)} 个模型:")
            for model in models:
                print(f"   - {model['name']}")
            
            return True
        else:
            print("❌ Ollama连接失败")
            return False
            
    except Exception as e:
        print(f"❌ 测试Ollama连接时出错: {e}")
        traceback.print_exc()
        return False


def test_milvus_connection(bot):
    """测试Milvus连接"""
    if not bot:
        return False
        
    print("\n🗄️ 测试Milvus连接...")
    
    try:
        stats = bot.search_tool.get_collection_stats()
        print(f"✅ Milvus连接正常")
        print(f"   集合名称: {stats['collection_name']}")
        print(f"   数据条数: {stats['total_entities']:,}")
        return True
        
    except Exception as e:
        print(f"❌ Milvus连接失败: {e}")
        traceback.print_exc()
        return False


def test_file_search(bot):
    """测试文件搜索功能"""
    if not bot:
        return False
        
    print("\n🔍 测试文件搜索功能...")
    
    try:
        query = "python"
        print(f"搜索关键词: '{query}'")
        
        results = bot.search_files(query, top_k=3)
        
        if results:
            print(f"✅ 搜索成功，找到 {len(results)} 个文件:")
            for i, result in enumerate(results, 1):
                print(f"   {i}. {result['file_name']} (相似度: {result['similarity_score']:.3f})")
            return True
        else:
            print("⚠️ 搜索结果为空")
            return False
            
    except Exception as e:
        print(f"❌ 文件搜索失败: {e}")
        traceback.print_exc()
        return False


def test_context_building(bot):
    """测试上下文构建"""
    if not bot:
        return False
        
    print("\n📝 测试上下文构建...")
    
    try:
        query = "python文件"
        files = bot.search_files(query, top_k=2)
        
        if files:
            context = bot.build_context(query, files)
            print(f"✅ 上下文构建成功")
            print(f"   上下文长度: {len(context)} 字符")
            print(f"   包含文件: {len(files)} 个")
            return True
        else:
            print("⚠️ 没有文件用于构建上下文")
            return False
            
    except Exception as e:
        print(f"❌ 上下文构建失败: {e}")
        traceback.print_exc()
        return False


def test_response_generation(bot):
    """测试响应生成"""
    if not bot:
        return False
        
    print("\n🤖 测试响应生成...")
    
    try:
        query = "测试问题"
        context = "用户问题: 测试问题\n相关文件信息：\n文件1: test.py\n路径: /test/test.py\n类型: .py"
        
        print("生成响应中...")
        response = bot.generate_response(query, context)
        
        if response:
            print(f"✅ 响应生成成功")
            print(f"   响应长度: {len(response)} 字符")
            print(f"   响应预览: {response[:100]}...")
            return True
        else:
            print("❌ 响应为空")
            return False
            
    except Exception as e:
        print(f"❌ 响应生成失败: {e}")
        traceback.print_exc()
        return False


def test_chat_method(bot):
    """测试chat方法"""
    if not bot:
        return False
        
    print("\n💬 测试chat方法...")
    
    try:
        question = "有哪些Python文件？"
        print(f"问题: '{question}'")
        
        print("调用chat方法...")
        result = bot.chat(question)
        
        print(f"chat方法返回类型: {type(result)}")
        print(f"chat方法返回值: {result}")
        
        if result is None:
            print("❌ chat方法返回None")
            return False
        elif isinstance(result, tuple) and len(result) == 2:
            response, sources = result
            print(f"✅ chat方法返回正确的元组")
            print(f"   响应类型: {type(response)}")
            print(f"   响应长度: {len(response) if response else 0}")
            print(f"   源文件数量: {len(sources) if sources else 0}")
            return True
        else:
            print(f"❌ chat方法返回格式异常: {type(result)}")
            return False
            
    except Exception as e:
        print(f"❌ chat方法测试失败: {e}")
        traceback.print_exc()
        return False


def main():
    """主测试函数"""
    print("🔍 RAG系统诊断测试")
    print("=" * 60)
    
    # 1. 测试初始化
    bot = test_rag_initialization()
    if not bot:
        print("\n❌ 初始化失败，无法继续测试")
        return
    
    # 2. 测试Ollama连接
    ollama_ok = test_ollama_connection(bot)
    
    # 3. 测试Milvus连接
    milvus_ok = test_milvus_connection(bot)
    
    # 4. 测试文件搜索
    search_ok = test_file_search(bot)
    
    # 5. 测试上下文构建
    context_ok = test_context_building(bot)
    
    # 6. 测试响应生成
    response_ok = test_response_generation(bot)
    
    # 7. 测试chat方法
    chat_ok = test_chat_method(bot)
    
    # 总结
    print("\n" + "=" * 60)
    print("📊 测试结果总结:")
    print(f"   初始化: {'✅' if bot else '❌'}")
    print(f"   Ollama连接: {'✅' if ollama_ok else '❌'}")
    print(f"   Milvus连接: {'✅' if milvus_ok else '❌'}")
    print(f"   文件搜索: {'✅' if search_ok else '❌'}")
    print(f"   上下文构建: {'✅' if context_ok else '❌'}")
    print(f"   响应生成: {'✅' if response_ok else '❌'}")
    print(f"   Chat方法: {'✅' if chat_ok else '❌'}")
    
    if all([bot, ollama_ok, milvus_ok, search_ok, context_ok, response_ok, chat_ok]):
        print("\n🎉 所有测试通过，RAG系统正常!")
    else:
        print("\n⚠️ 部分测试失败，需要检查相关组件")


if __name__ == "__main__":
    main()